package com.htd.base.common.filter;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import jakarta.servlet.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.annotation.Order;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.io.IOException;

@Slf4j
public class DataSourceContextFilter implements Filter {
    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        try {
            // 获取当前数据源
            String currentDataSource = DynamicDataSourceContextHolder.peek();
            if (currentDataSource != null) {
                // 保存到请求上下文中
                RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
                if (requestAttributes != null) {
                    requestAttributes.setAttribute("currentDataSource", currentDataSource, RequestAttributes.SCOPE_REQUEST);
                    log.debug("[Trans. plugin] Save datasource to request context: {}", currentDataSource);
                }
            }
            chain.doFilter(request, response);
        } finally {
            // 清理请求上下文中的数据源信息
            RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
            if (requestAttributes != null) {
                requestAttributes.removeAttribute("currentDataSource", RequestAttributes.SCOPE_REQUEST);
            }
        }
    }
}
